//! 多线程处理模块
//! 
//! 提供多线程并行数据生成的支持，提升数据生成的速度和效率

use rayon::prelude::*;
use std::sync::{Arc, Mutex};
use std::collections::HashMap;
use serde_json::Value;
use crate::error::{DataForgeError, Result};

pub mod thread_pool;
pub mod safety;

pub use thread_pool::*;
pub use safety::*;

/// 并行数据生成器
pub struct ParallelGenerator {
    /// 线程池大小
    thread_count: usize,
    /// 批次大小
    batch_size: usize,
    /// 生成器函数
    generators: HashMap<String, Arc<dyn Fn() -> Value + Send + Sync>>,
}

impl ParallelGenerator {
    /// 创建新的并行生成器
    pub fn new(thread_count: usize, batch_size: usize) -> Self {
        Self {
            thread_count,
            batch_size,
            generators: HashMap::new(),
        }
    }

    /// 注册生成器
    pub fn register<F>(&mut self, name: String, generator: F)
    where
        F: Fn() -> Value + Send + Sync + 'static,
    {
        self.generators.insert(name, Arc::new(generator));
    }

    /// 并行生成数据
    pub fn generate_parallel(&self, generator_name: &str, count: usize) -> Result<Vec<Value>> {
        let generator = self.generators.get(generator_name)
            .ok_or_else(|| DataForgeError::validation(&format!("Generator '{}' not found", generator_name)))?;

        // 配置rayon线程池
        let pool = rayon::ThreadPoolBuilder::new()
            .num_threads(self.thread_count)
            .build()
            .map_err(|e| DataForgeError::generator(&format!("Failed to create thread pool: {}", e)))?;

        let results = Arc::new(Mutex::new(Vec::with_capacity(count)));
        let generator = Arc::clone(generator);

        pool.install(|| {
            (0..count).into_par_iter().for_each(|_| {
                let value = generator();
                results.lock().unwrap().push(value);
            });
        });

        let results = Arc::try_unwrap(results)
            .map_err(|_| DataForgeError::generator("Failed to unwrap results"))?
            .into_inner()
            .map_err(|_| DataForgeError::generator("Failed to acquire mutex"))?;

        Ok(results)
    }

    /// 批量并行生成数据
    pub fn generate_batched(&self, generator_name: &str, total_count: usize) -> Result<Vec<Value>> {
        let generator = self.generators.get(generator_name)
            .ok_or_else(|| DataForgeError::validation(&format!("Generator '{}' not found", generator_name)))?;

        let batches: Vec<usize> = (0..total_count)
            .step_by(self.batch_size)
            .map(|start| {
                let end = (start + self.batch_size).min(total_count);
                end - start
            })
            .collect();

        let generator = Arc::clone(generator);
        let results: Result<Vec<Vec<Value>>> = batches
            .into_par_iter()
            .map(|batch_size| {
                let mut batch_results = Vec::with_capacity(batch_size);
                for _ in 0..batch_size {
                    batch_results.push(generator());
                }
                Ok(batch_results)
            })
            .collect();

        let results = results?;
        Ok(results.into_iter().flatten().collect())
    }

    /// 获取线程数
    pub fn thread_count(&self) -> usize {
        self.thread_count
    }

    /// 获取批次大小
    pub fn batch_size(&self) -> usize {
        self.batch_size
    }
}

impl Default for ParallelGenerator {
    fn default() -> Self {
        Self::new(num_cpus::get(), 1000)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_parallel_generator_creation() {
        let generator = ParallelGenerator::new(4, 100);
        assert_eq!(generator.thread_count(), 4);
        assert_eq!(generator.batch_size(), 100);
    }

    #[test]
    fn test_register_generator() {
        let mut generator = ParallelGenerator::new(2, 50);
        generator.register("test".to_string(), || Value::String("test".to_string()));
        
        let results = generator.generate_parallel("test", 10);
        assert!(results.is_ok());
        assert_eq!(results.unwrap().len(), 10);
    }

    #[test]
    fn test_batched_generation() {
        let mut generator = ParallelGenerator::new(2, 3);
        generator.register("counter".to_string(), || {
            Value::Number(serde_json::Number::from(42))
        });
        
        let results = generator.generate_batched("counter", 10);
        assert!(results.is_ok());
        assert_eq!(results.unwrap().len(), 10);
    }
}